import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import pywt
from config import Config
from collections import defaultdict
import warnings
warnings.filterwarnings("ignore")

# 熵函数
def entropy(data):
    len_slip = 0.01
    defdic = defaultdict(int)
    for i in data:
        defdic[int(round(i / len_slip))] += 1
    E=0
    num_x = len(data)
    for i in defdic:
        E = E - defdic[i]/num_x*np.log(defdic[i]/num_x)
    return E

# 时域特征提取
def time_domain_feature(Config):
    path_data = Config.path_data 
    path_data = path_data.replace('\\','_')+'.csv'
    data_bearing = pd.read_csv(path_data, header=None)
    print(data_bearing)
    print('采样个数:',data_bearing.shape[0])
    plt.plot(data_bearing.iloc[0])
    plt.title('单个采样数据')
    plt.show()

    Peak = []
    Peak2peak = []
    RMS = []
    Skewness = []
    Crest_factor = []
    Kurtosis_values = []
    Impulse_factor = []
    Form_factor = []
    Clearance_factor = []
    Kurtosis = []
    Entropy = []

    for i in range(data_bearing.shape[0]):
        data = data_bearing.iloc[i].values
        Peak.append(max(np.abs(data)))
        Peak2peak.append(max(data)-min(data))
        RMS.append(np.mean(data**2)**0.5)
        Skewness.append(np.mean(data**3))

        Crest_factor.append(Peak[-1]/RMS[-1])
        Kurtosis_values.append( np.mean(np.abs(data-np.mean(data))**4/RMS[-1]**4) )
        Impulse_factor.append(max(data)/np.abs(np.mean(data)))
        Form_factor.append(RMS[-1]/np.abs(np.mean(data)))
        Clearance_factor.append(Peak[-1]/np.abs(np.mean(np.abs(data)**0.5))**2)
        Kurtosis.append(np.mean((data-np.mean(data))**4)/np.mean((data-np.mean(data))**2)**2)
        Entropy.append(entropy(data))
    time_domain_data=[Peak,Peak2peak,RMS,Skewness,Crest_factor,Kurtosis_values,Impulse_factor,Form_factor,Clearance_factor,Kurtosis,Entropy]

    '''
    plt.figure(figsize=(18,9))
    plt.subplot(3,4,1)
    plt.plot(Peak,label='Peak')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,2)
    plt.plot(Peak2peak,label='Peak2peak')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,3)
    plt.plot(RMS,label='RMS')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,4)
    plt.plot(Skewness,label='skewness')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,5)
    plt.plot(Crest_factor,label='Crest_factor')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,6)
    plt.plot(Kurtosis_values,label='Kurtosis_values')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,7)
    plt.plot(Impulse_factor,label='Impulse_factor')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,8)
    plt.plot(Form_factor,label='Form_factor')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,9)
    plt.plot(Clearance_factor,label='Clearance_factor')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,10)
    plt.plot(Kurtosis,label='Kurtosis')
    plt.legend()
    # plt.show()
    plt.subplot(3,4,11)
    plt.plot(Entropy,label='Entropy')
    plt.legend()
    plt.subplots_adjust(top=0.98,bottom=0.05,left=0.05,right=0.98) #,hspace=0.08,wspace=0.08
    plt.savefig('Feature.png')
    plt.show()
    '''
    return np.array(time_domain_data).T


# 傅里叶变换（取一半）
def fft1(xx):
    t=np.linspace(0, 1.0, len(xx))
    f = np.arange(len(xx)/2+1, dtype=complex)
    for index in range(len(f)):
        f[index]=complex(np.sum(np.cos(2*np.pi*index*t)*xx), -np.sum(np.sin(2*np.pi*index*t)*xx))
    return f

# 时域特征提取
def frequency_domain_feature(Config):
    path_data = Config.path_data 
    path_data = path_data.replace('\\','_')+'.csv'
    data_bearing = pd.read_csv(path_data, header=None)
    print(data_bearing)
    print('采样个数:',data_bearing.shape[0])
    plt.plot(data_bearing.iloc[0])
    plt.title('单个采样数据')
    plt.show()

    sampling_rate=25.6*1000   # 采样频率为25.6kHz
    frequency_data =[]
    # plt.ion()
    # plt.figure(figsize=(12,5))
    
    for i in range(data_bearing.shape[0]):
        data = data_bearing.iloc[i].values
        x1 =data
        xf=fft1(x1)/len(x1)  # 求均值
        freqs = np.linspace(0, sampling_rate/2, int(len(x1)/2+1))  # 最大频率为采样频率的一半
        # plt.cla() # 清空当前图
        # plt.plot(freqs,2*np.abs(xf),'r-')  # 频率幅值乘以2
        # plt.xlabel("Frequency(Hz)")
        # plt.ylabel("Amplitude($m$)")
        # plt.title("Amplitude-Frequency curve")
        # plt.ylim(0,0.15)
        # plt.pause(0.01)
        # 对频谱进行划分,求各个划分段的频谱幅值
        frequency_data.append([np.sum(2*np.abs(xf)[int(len(xf)/4*i):int(len(xf)/4*(i+1))]) for i in range(4)])
    plt.show()
    return np.array(frequency_data)


def wavelet_feature(Config):   # reference:https://blog.csdn.net/qq_41978536/article/details/89337436
    path_data = Config.path_data
    path_data = path_data.replace('\\','_')+'.csv'
    data_bearing = pd.read_csv(path_data, header=None)
    print(data_bearing)
    print('采样个数:',data_bearing.shape[0])
    plt.plot(data_bearing.iloc[0])
    plt.title('单个采样数据')
    plt.show()

    # plt.ion()
    # fig = plt.figure(figsize=(16,9))
    wavelet_data = []
    for i in range(data_bearing.shape[0]):
        data = data_bearing.iloc[i].values
        wp = pywt.WaveletPacket(data=data, wavelet='db5', mode='symmetric')
        # print('wp.maxlevel:',wp.maxlevel)
        node_level = [node.path for node in wp.get_level(4, 'natural')] # 按自然顺序得到特定层上的所有节点
        # plt.clf()
        # fig.suptitle('小波包分解')
        w_data = []
        E_sum = 0
        for n,c in enumerate(node_level):
            w_data.append(sum(wp[c].data**2))
            E_sum+=sum(wp[c].data**2)
        #     plt.subplot(4,4,n+1)
        #     plt.plot(wp[c].data,label=c)
        #     plt.legend()
        # plt.pause(0.01)
        wavelet_data.append(np.array(w_data)/E_sum)
    return np.array(wavelet_data)


if __name__ =='__main__':
    time_domain_data = time_domain_feature(Config)
    frequency_data = frequency_domain_feature(Config)
    wavelet_data = wavelet_feature(Config)
    feature_data = np.hstack((time_domain_data, frequency_data, wavelet_data))
    np.savetxt('new.csv', feature_data, delimiter = ',')
